#!/usr/bin/env python
#
# Copyright 2011 (C) David I. Lehn <dil@lehn.org>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import logging
import random
from google.appengine.ext import db
from google.appengine.ext.deferred import deferred

from sparserange import SparseRange


# default initial shard count
DEFAULT_SHARD_COUNT = 1
# default max allocator shards
DEFAULT_SHARD_MAX = 20
# default allocated ids
DEFAULT_ALLOCATED_IDS = "0"
# default IdAllocator size at which to schedule a refill task
DEFAULT_MIN_COUNT = 100
# default number of ids to refill an IdAllocator
DEFAULT_REFILL_COUNT = 500


class IdAllocatorConfig(db.Model):
    """Tracks the number of shards for each named allocator."""
    v = db.IntegerProperty(default=1)
    shard_count = db.IntegerProperty(default=DEFAULT_SHARD_COUNT)
    shard_max = db.IntegerProperty(default=DEFAULT_SHARD_MAX)
    allocated_ids = db.StringProperty(default=DEFAULT_ALLOCATED_IDS)
    min_count = db.IntegerProperty(default=DEFAULT_MIN_COUNT)
    refill_count = db.IntegerProperty(default=DEFAULT_REFILL_COUNT)

    def allocate_ids(self, count=0):
        """Allocate a SparseRange of at least count ids."""
        def txn():
            sr = SparseRange(self.allocated_ids)
            ids = sr.add_count(count + self.refill_count)
            self.allocated_ids = str(sr)
            self.put()
            return SparseRange(ids)
        return db.run_in_transaction(txn)


def refill_id_allocator(allocator, config):
    """Refill an IdAllocator."""
    sr = SparseRange(allocator.available_ids)
    if len(sr) < config.min_count:
        ids = config.allocate_ids()
        if ids is not None:
            def txn():
                sr = SparseRange(allocator.available_ids)
                sr.add_range(ids)
                allocator.available_ids = str(sr)
                allocator.put()
            db.run_in_transaction(txn)


def refill_id_allocator_task(k):
    """Refill an IdAllocator specified by key."""
    allocator = IdAllocator.get(k)
    config = allocator.config
    refill_id_allocator(allocator, config)


class IdAllocator(db.Model):
    """Shard for a named allocator."""
    v = db.IntegerProperty(default=1)
    config = db.ReferenceProperty(IdAllocatorConfig)
    available_ids = db.StringProperty(required=True)

    def get_ids(self, config, count=1):
        """Get a SparseRange of count ids.

        This method should only be called in a transaction.
        """
        # Note: config param due to entity group issues in transactionss
        sr = SparseRange(self.available_ids)
        ids = None
        try:
            ids = sr.remove_count(count)
        except IndexError:
            pass
        if ids is not None:
            self.available_ids = str(sr)
            self.put()
            # schedule refill if needed
            if len(sr) < config.min_count:
                deferred.defer(refill_id_allocator_task, self.key())
        return ids


def allocate_ids(name, count=1):
    """Allocate IDs.

    Parameters:
      name - The name of the allocator.
      count - The number of ids to return.
    """
    config = IdAllocatorConfig.get_or_insert(name)

    ids = None
    attempts = 10
    new_shared_attempt = 9
    while ids is None and attempts > 0:
        attempts -= 1

        # get next potential allocator info
        # (id 0 reserved)
        index = random.randint(1, config.shard_count)
        shard_name = '%s[%d]' % (name, index)
        allocator = None
        # try to get allocator and new ids
        def txn():
            allocator = IdAllocator.get_by_key_name(shard_name)
            if allocator is None:
                return (None, None)
            return allocator, allocator.get_ids(config, count)
        allocator, ids = db.run_in_transaction(txn)

        # need a new allocator with new range
        if allocator is None:
            # get new range (uses a config transaction)
            # this range may be lost if futher processing fails
            ia_ids = config.allocate_ids()

            # try to get allocator again, else create w/ ids
            # new ids lost if allocator already created
            def txn():
                allocator = IdAllocator.get_by_key_name(shard_name)
                if allocator is None:
                    allocator = IdAllocator(
                            key_name=shard_name, config=config,
                            available_ids=str(ia_ids))
                    allocator.put()
                else:
                    logging.info("allocate_ids lost ids '%s'" % str(ia_ids))
                return allocator.get_ids(config, count)
            ids = db.run_in_transaction(txn)
        elif ids is None:
            # refill this allocator and try to use it
            refill_id_allocator(allocator, config)
            def txn():
                return allocator.get_ids(config, count)
            ids = db.run_in_transaction(txn)

        # high failure rate, check if shard count should be increased
        if attempts == new_shared_attempt and \
                config.shard_count < config.shard_max:
            def txn():
                if config.shard_count < config.shard_max:
                    config.shard_count += 1
                    config.put()
                    logging.info("increased shard count for '%s' to %d",
                            name, config.shard_count)
            db.run_in_transaction(txn)

    return ids

# TODO
# - full cache levels:
#   - datastore
#     - all valid data
#   - memcache
#   - in-memory: store a few for high speed access
#     - could write back if not used
